
/**
 * @file aicl_mdl_001.cpp
 *
 * Copyright (c) Huawei Technologies Co., Ltd. 2012-2018. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include <iostream>
#include <fstream>
#include <cstring>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>

#include <stdio.h>
#include <math.h>
#include <ctime>
#include <chrono>
#include <map>

#include "acl/acl.h"
#include "aicl.h"
#include "utils.h"

// sys
int32_t deviceId_ = 0;
uint32_t modelId;
aiclRet ret;
aiclmdlDesc *modelDesc;
size_t binDataSize = 0;
aiclmdlDataset *inputDataSet;
aiclDataBuffer *inputDataBuffer;
aiclmdlDataset *outputDataSet;
aiclDataBuffer *outputDataBuffer;
size_t outputDataSize = 0;
void *g_binHostData;
void *g_binDeviceData;
void *g_outputDeviceData;
void *g_outputHostData;

using namespace std;

bool ReadBinTotHost(string fileName)
{
    ifstream binFile(fileName, ifstream::binary);
    binFile.seekg(0, binFile.end);
    binDataSize = binFile.tellg();
    INFO_LOG("load bin file size %d", binDataSize);
    if (binDataSize == 0) {
        return false;
    }
    binFile.seekg(0, binFile.beg);
    aiclrtMallocHost(&g_binHostData, binDataSize);
    binFile.read((char*)g_binHostData, binDataSize);
    binFile.close();
    return true;
}

void UnloadBuffer()
{
    aiclrtFreeHost(g_binHostData);
    g_binHostData = nullptr;
    aiclrtFree(g_binDeviceData);
    g_binDeviceData = nullptr;
    aiclDestroyDataBuffer(inputDataBuffer);
    inputDataBuffer = nullptr;
    aiclmdlDestroyDataset(inputDataSet);
    inputDataSet = nullptr;

    aiclrtFreeHost(g_outputHostData);
    g_outputHostData = nullptr;
    aiclrtFree(g_outputDeviceData);
    g_outputDeviceData = nullptr;
    aiclDestroyDataBuffer(outputDataBuffer);
    outputDataBuffer = nullptr;
    aiclmdlDestroyDataset(outputDataSet);
    outputDataSet = nullptr;
}

int main()
{
    const char *workPath = getenv("OPENCANN_TPATH");
    if (workPath == NULL) {
        ERROR_LOG("OPENCANN_TPATH is not set");
        return OPENCANN_FAILED;
    }
    // 1. Init
    std::string workPath_s = workPath;
    std::string configPath = workPath_s + "/common/test_data/config/init_config.json";

    aiclInit(configPath.c_str());
    aiclrtSetDevice(deviceId_);

    //std::string omModelPath = workPath_s + "/common/test_data/model/aicl_mdl_conv.om";
    std::string omModelPath = workPath_s + "/common/test_data/model/resnet50.om";
    ret = aiclmdlLoadFromFile(omModelPath.c_str(), &modelId);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("load model from file failed, model file is %s, errorCode is %d", omModelPath.c_str(),
            static_cast<int32_t>(ret));
        return OPENCANN_FAILED;
    }

    // 2. load data bin
    string binPath = workPath_s + "/common/test_data/data/aicl_mdl_conv_input.bin";
    bool retBin = ReadBinTotHost(binPath);
    if (!retBin) {
        ERROR_LOG("load bin input file failed");
        return OPENCANN_FAILED;
    }
    ret = aiclrtMalloc(&g_binDeviceData, binDataSize, AICL_MEM_MALLOC_HUGE_FIRST);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("malloc %d failed", binDataSize);
        return OPENCANN_FAILED;
    }
    ret = aiclrtMemcpy(g_binDeviceData, binDataSize, g_binHostData, binDataSize, AICL_MEMCPY_HOST_TO_DEVICE);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("memcpy failed");
        return OPENCANN_FAILED;
    }

    // 3.prepare input desc
    inputDataSet = aiclmdlCreateDataset();
    inputDataBuffer = aiclCreateDataBuffer(g_binDeviceData, binDataSize);
    ret = aiclmdlAddDatasetBuffer(inputDataSet, inputDataBuffer);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("add input dataset buffer failed, errorCode is %d", static_cast<int32_t>(ret));
        (void)aiclDestroyDataBuffer(inputDataBuffer);
        inputDataBuffer = nullptr;
        return OPENCANN_FAILED;
    }
    INFO_LOG("create model input success");
   
    // 4.prepare output desc
    modelDesc =  aiclmdlCreateDesc();
    ret = aiclmdlGetDesc(modelDesc, modelId);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("get model description failed, modelId is %u, errorCode is %d", modelId, static_cast<int32_t>(ret));
        return OPENCANN_FAILED;
    }
    INFO_LOG("create model description success");
    
    outputDataSet = aiclmdlCreateDataset();
    outputDataSize = aiclmdlGetOutputSizeByIndex(modelDesc, 0);
    ret = aiclrtMalloc(&g_outputDeviceData, outputDataSize, AICL_MEM_MALLOC_HUGE_FIRST);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("memcpy failed");
        return OPENCANN_FAILED;
    }

    outputDataBuffer = aiclCreateDataBuffer(g_outputDeviceData, outputDataSize);
    ret = aiclmdlAddDatasetBuffer(outputDataSet, outputDataBuffer);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("can't add data buffer, create output failed, errorCode is %d", static_cast<int32_t>(ret));
        (void)aiclrtFree(g_outputDeviceData);
        (void)aiclDestroyDataBuffer(outputDataBuffer);
        return OPENCANN_FAILED;
    }
    INFO_LOG("create model output success");

    // 5.model execute
    ret = aiclmdlExecute(modelId, inputDataSet, outputDataSet);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("execute model failed, errorCode is %d", static_cast<int32_t>(ret));
        return OPENCANN_FAILED;
    }
    INFO_LOG("model execute success");

    // 6.unload
    aiclmdlDestroyDesc(modelDesc);
    ret = aiclmdlUnload(modelId);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("unload model failed, errorCode is %d", static_cast<int32_t>(ret));
        return OPENCANN_FAILED;
    }
    INFO_LOG("model unload success");

    UnloadBuffer();
    aiclrtResetDevice(deviceId_);
    aiclFinalize();


    return OPENCANN_SUCCESS;
}

